from typing import Dict, List, Set
from collections import defaultdict, deque

__all__ = ["DAG"]


class DAG:
    """
    有向无环图 (DAG) 实现，用于管理任务或节点之间的依赖关系。

    该类支持：
    - 动态添加节点及其依赖。
    - 从字典配置加载图结构（可自定义依赖字段名）。
    - 计算目标节点及其所有依赖的分层拓扑排序（执行计划）。
    - 以可读格式输出图结构。

    Attributes:
        graph (Dict[str, Set[str]]): 正向图，graph[predecessor] = {successor, ...}。
        reverse_graph (Dict[str, Set[str]]): 反向图，reverse_graph[successor] = {predecessor, ...}。
        nodes (Set[str]): 所有节点的集合。
        depends_on (str): 在加载配置字典时，用于查找依赖列表的键名，默认为 'deps'。
    """

    def __init__(self, graph_def: Dict[str, Dict[str, List[str]]] = None, deps: str = 'deps') -> None:
        """
        初始化 DAG 实例。

        Args:
            deps (str): 指定配置字典中表示依赖关系的键名。默认为 'deps'。
            graph_def (Dict[str, Dict[str, List[str]]], optional): 可选的初始图定义字典。
                                                                  如果提供，则立即加载。
        """
        self.graph: Dict[str, Set[str]] = defaultdict(set)
        self.reverse_graph: Dict[str, Set[str]] = defaultdict(set)
        self.nodes: Set[str] = set()
        self.depends_on = deps

        # 如果提供了 graph_def，则立即加载
        if graph_def is not None:
            self.load(graph_def)  # 复用现有的 load 方法

    def __str__(self) -> str:
        """
        返回 DAG 的字符串表示，用于 print() 和 str()。
        默认显示整个图的拓扑结构。
        """
        # 复用现有的 dumps 逻辑，但默认不显示 reverse_graph raw
        return self.dumps(targets=None, sort=True, show_reverse=False)

    def add_dep(self, node: str, dep: List[str]) -> None:
        """
        向图中添加一个节点及其依赖项。

        Args:
            node (str): 要添加的节点名称。
            dep (List[str]): 该节点所依赖的节点名称列表。
        """
        self.nodes.add(node)
        for d in dep:
            self.graph[d].add(node)
            self.reverse_graph[node].add(d)
            self.nodes.add(d)

    def load(self, graph_def: Dict[str, Dict[str, List[str]]]) -> None:
        """
        从字典定义中批量加载 DAG 结构。

        字典的每个键是一个节点名，值是一个包含依赖信息的字典。
        依赖列表通过实例化时指定的 `depends_on` 键名获取。

        Args:
            graph_def (Dict[str, Dict[str, List[str]]]): 描述 DAG 的字典。
                示例: {"task_a": {"deps": ["task_b"]}, "task_b": {"deps": []}}

        Raises:
            KeyError: 如果配置字典的值不是字典或缺少依赖键，`get` 方法会返回空列表，通常不会抛出 KeyError。
                      此处主要依赖 `add_dep` 的逻辑。
        """
        for node, conf in graph_def.items():
            # 假设 conf 是一个字典，从中获取依赖列表
            dependencies = conf.get(self.depends_on, [])
            self.add_dep(node, dependencies)

    def task_chain(self, targets: Set[str]) -> List[List[str]]:
        """
        计算从目标节点反向追溯的所有依赖节点的分层拓扑排序。

        结果是一个列表的列表，外层列表的每个元素代表一个执行层。
        同一层内的任务可以并行执行，下一层的任务依赖于上一层所有任务的完成。

        Args:
            targets (Set[str]): 目标节点集合，计算这些节点及其所有祖先依赖的执行计划。

        Returns:
            List[List[str]]: 分层的拓扑排序列表。每层内的节点已按字典序排序。

        Raises:
            KeyError: 如果 `targets` 中包含图中未定义的节点。
            ValueError: 如果子图中存在循环依赖。
        """
        unknown = targets - self.nodes
        if unknown:
            raise KeyError(f"未知节点: {unknown}")

        # 1. 收集目标节点及其所有祖先节点（反向遍历 reverse_graph）
        visited: Set[str] = set()
        q = deque(targets)
        while q:
            node = q.popleft()
            if node in visited:
                continue
            visited.add(node)
            # 将当前节点的所有依赖（父节点）加入队列
            q.extend(self.reverse_graph.get(node, []))

        # 2. 为收集到的子图构建入度表（基于 reverse_graph 计算每个节点在子图中的入度）
        in_degree: Dict[str, int] = {n: 0 for n in visited}
        for n in visited:
            for d in self.reverse_graph.get(n, []):
                if d in visited:  # 确保依赖也在子图内
                    in_degree[n] += 1

        # 3. Kahn's Algorithm 进行分层拓扑排序
        layers: List[List[str]] = []
        # 当前层：子图中所有入度为0的节点
        cur: List[str] = [n for n in visited if in_degree[n] == 0]
        processed = 0
        while cur:
            layers.append(sorted(cur))  # 对当前层节点排序以保证输出一致性
            processed += len(cur)
            nxt: List[str] = []
            for n in cur:
                # 遍历当前节点的所有后继（子节点）
                for child in self.graph.get(n, []):
                    if child in visited:  # 确保子节点也在子图内
                        in_degree[child] -= 1
                        if in_degree[child] == 0:
                            nxt.append(child)
            cur = nxt

        # 检查是否所有节点都被处理（无环）
        if processed != len(visited):
            raise ValueError("图中存在环，无法分层")
        return layers

    def dumps(self, targets: Set[str] = None, sort: bool = True, show_reverse: bool = False) -> str:
        """
        生成 DAG 结构的可读字符串表示。

        Args:
            targets (Set[str], optional): 如果指定，只显示这些目标节点及其依赖构成的子图。
                                          默认为 None，显示整个图。
            sort (bool): 是否对每层内的节点进行排序。默认为 True。
            show_reverse (bool): 是否在末尾打印原始的 reverse_graph。默认为 False。

        Returns:
            str: 多行字符串，展示图的分层结构和依赖关系。
        """
        lines = ["DAG Graph (topologically sorted):"]

        # 确定要显示的节点集
        display_nodes: List[str]
        if targets is None:
            display_nodes = list(self.nodes)
        else:
            display_nodes = list(targets)

        # 计算拓扑分层
        try:
            layers = self.task_chain(set(display_nodes))
        except (KeyError, ValueError) as e:
            lines.append(f"[!] 拓扑排序失败: {e}")
            # 回退：直接显示给定的节点（不排序）
            layers = [sorted(display_nodes)]

        # 构建输出行
        for i, layer in enumerate(layers):
            lines.append(f"  Layer {i}:")
            for node in (sorted(layer) if sort else layer):
                deps = sorted(self.reverse_graph[node])  # 依赖项（父节点）
                children = sorted(self.graph[node])  # 被谁使用（子节点）
                lines.append(f"    {node}: {deps}")
                # lines.append(f"      depends on: {deps}")
                # lines.append(f"      used by:    {children}")

        # 可选：显示原始 reverse_graph
        if show_reverse:
            lines.append("\n[reverse_graph raw]:")
            for node in sorted(self.reverse_graph):
                lines.append(f"  {node} <- {sorted(self.reverse_graph[node])}")

        return "\n".join(lines)



